#from openai import OpenAI
from zhipuai import ZhipuAI
from dotenv import load_dotenv
import os
import json

load_dotenv()

def send_messages(messages):
    response = client.chat.completions.create(
        model="glm-4",
        messages=messages,
        tools=tools
    )
    print("messages: ", messages)
    return response.choices[0].message

# client = OpenAI(
#     api_key=os.getenv("DEEPSEEK_API_KEY"),
#     base_url="https://api.deepseek.com",
# )
client = ZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY")) # ZHIPUAI_API_KEY

tools = [
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "Get weather of an location, the user shoud supply a location first",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "The city and state, e.g. San Francisco, CA",
                    }
                },
                "required": ["location"]
            },
        }
    },
]

messages = [{"role": "user", "content": "How's the weather in Hangzhou?"}]
print(f"User>\t {messages[0]['content']}")

message = send_messages(messages)
print("\nModel's Function Call Response:")
print(f"Content: {message.content}")
print(f"Tool Calls: {json.dumps(message.tool_calls[0].function.model_dump(), indent=2)}")
print("-" * 200)
# tool_call = message.tool_calls[0]
# messages.append(message)
# messages.append({"role": "tool", "tool_call_id": tool_call.id, "content": "24℃"})
# message = send_messages(messages)
# print(f"\nModel's Final Response: {message.content}")

# 智谱AI返回结构处理
if hasattr(message, 'tool_calls') and message.tool_calls:
    tool_call = message.tool_calls[0]
    print(f"Content: {message.content}")
    print(f"Tool Calls: {json.dumps({'name': tool_call.function.name,'arguments': tool_call.function.arguments}, indent=2)}")

    # 添加到消息历史
    messages.append({
        "role": "assistant",
        "content": message.content,
        "tool_calls": [{
            "id": tool_call.id,
            "function": {
                "name": tool_call.function.name,
                "arguments": tool_call.function.arguments
            },
            "type": "function"
        }]
    })

    # 模拟工具返回
    messages.append({
        "role": "tool",
        "content": "24℃, Sunny",
        "tool_call_id": tool_call.id
    })

    # 第二次调用
    final_response = send_messages(messages)
    print(f"\nModel's Final Response: {final_response.content}")
else:
    print("No tool calls were returned.")